import os
import numpy as np
import nibabel as nib
import torch
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split
import pickle

class MRIDataset(Dataset):
    def __init__(self, root_dir, ed_es_path, split='train', transform=None, val_size=0.2, random_state=42):
        """
        Args:
            root_dir (string): Directory with all the folders containing image and segmentation map pairs.
            ed_es_path (string): Path to the pickle file containing ED and ES time steps.
            split (string): 'train' for training data or 'val' for validation data.
            transform (callable, optional): Optional transform to be applied on a sample.
            val_size (float): Fraction of data to be used for validation.
            random_state (int): Random state for reproducibility.
        """
        self.root_dir = root_dir
        self.transform = transform
        self.split = split

        all_samples = self._load_samples()
        train_samples, val_samples = train_test_split(all_samples,
                                                      test_size=val_size,
                                                      random_state=random_state)
        self.samples = train_samples if self.split == 'train' else val_samples

        with open(ed_es_path, 'rb') as f:
            self.ed_es_dict = pickle.load(f)

    def _load_samples(self):
        samples = []
        for folder_name in os.listdir(self.root_dir):
            folder_path = os.path.join(self.root_dir, folder_name)
            if os.path.isdir(folder_path):
                image_path = os.path.join(folder_path, 'sa_cropped.nii.gz')
                if os.path.exists(image_path):
                    ed_es_info = self.ed_es_dict.get(folder_name, {})
                    samples.append((image_path, ed_es_info))
        return samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        image_path, ed_es_info = self.samples[idx]

        img_nii = nib.load(image_path, mmap_mode='r')
        middle_slice_index = img_nii.shape[2] // 2
        image = img_nii.dataobj[:, :, middle_slice_index]
        image = torch.tensor(image, dtype=torch.float)

        ed_time_step = torch.tensor(ed_es_info['ED_time_step'], dtype=torch.int8)
        es_time_step = torch.tensor(ed_es_info['ES_time_step'], dtype=torch.int8)
    
        if self.transform:
            image = self.transform(image)

        return image, ed_time_step, es_time_step
